#!/bin/bash

# 定义虚拟机名称数组
VMs=("hadoop-app" "hadoop-node1" "hadoop-node2" "hadoop-node3")

# 配置参数
SLEEP_TIME=15  # 等待时间（秒），可以根据需要调整

# 颜色定义，用于美化输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印带颜色的信息函数
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}
print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}
print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}
print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示所有VM状态的函数
show_status() {
    print_info "当前 Hadoop 虚拟机状态："
    local all_running=true
    local any_running=false
    
    for vm in "${VMs[@]}"; do
        state=$(virsh domstate "$vm" 2>/dev/null || echo "未找到")
        if [[ "$state" == "running" ]]; then
            echo -e "  $vm: ${GREEN}运行中${NC}"
            any_running=true
        elif [[ "$state" == "shut off" ]]; then
            echo -e "  $vm: ${YELLOW}已关闭${NC}"
            all_running=false
        else
            echo -e "  $vm: ${RED}${state}${NC}"
            all_running=false
        fi
    done
    
    # 返回状态码，便于其他函数判断
    if $all_running && $any_running; then
        return 0  # 所有虚拟机都在运行
    elif $any_running; then
        return 1  # 部分虚拟机在运行
    else
        return 2  # 没有虚拟机在运行
    fi
}

# 等待并检查状态的函数
wait_and_check_status() {
    local wait_time=$1
    local operation=$2
    
    print_info "等待 ${wait_time} 秒让虚拟机${operation}完成..."
    
    # 显示进度条（可选）
    echo -n "等待中 "
    for i in $(seq 1 $wait_time); do
        echo -n "."
        sleep 1
    done
    echo ""
    
    print_info "等待完成，检查最终状态："
    show_status
}

# 检查是否所有虚拟机都启动成功
check_all_running() {
    local all_ok=true
    for vm in "${VMs[@]}"; do
        if ! virsh domstate "$vm" | grep -q "running"; then
            print_error "$vm 启动失败或仍在启动中"
            all_ok=false
        fi
    done
    
    if $all_ok; then
        print_success "所有虚拟机都已成功启动并运行！"
        return 0
    else
        print_warning "部分虚拟机可能尚未就绪，请稍后再次检查状态"
        return 1
    fi
}

# 启动所有VM的函数
start_vms() {
    print_info "正在启动 Hadoop 虚拟机..."
    local started_any=false
    
    for vm in "${VMs[@]}"; do
        if virsh domstate "$vm" | grep -q "shut off"; then
            print_info "启动 $vm..."
            virsh start "$vm"
            if [ $? -eq 0 ]; then
                print_success "$vm 启动命令执行成功"
                started_any=true
            else
                print_error "$vm 启动失败"
            fi
        else
            print_warning "$vm 不在关闭状态，跳过"
        fi
    done
    
    if $started_any; then
        print_success "虚拟机启动命令已发送"
        echo ""
        
        # 等待一段时间后检查状态
        wait_and_check_status $SLEEP_TIME "启动"
        
        # 最终检查所有虚拟机是否都运行
        check_all_running
    else
        print_info "没有需要启动的虚拟机"
        show_status
    fi
}

# 关闭所有VM的函数
shutdown_vms() {
    print_info "正在关闭 Hadoop 虚拟机..."
    local shutdown_any=false
    
    for vm in "${VMs[@]}"; do
        if virsh domstate "$vm" | grep -q "running"; then
            print_info "关闭 $vm..."
            virsh shutdown "$vm" # 发送ACPI关机信号，优雅关机
            if [ $? -eq 0 ]; then
                print_success "$vm 关机命令执行成功"
                shutdown_any=true
            else
                print_error "$vm 关机失败，尝试强制关闭..."
                virsh destroy "$vm" # 强制关机（相当于断电）
            fi
        else
            print_warning "$vm 不在运行状态，跳过"
        fi
    done
    
    if $shutdown_any; then
        print_success "虚拟机关机命令已发送"
        echo ""
        
        # 等待一段时间后检查状态
        wait_and_check_status $SLEEP_TIME "关闭"
    else
        print_info "没有需要关闭的虚拟机"
        show_status
    fi
}

# 重启所有VM的函数
restart_vms() {
    print_info "正在重启 Hadoop 虚拟机..."
    local restarted_any=false
    
    for vm in "${VMs[@]}"; do
        if virsh domstate "$vm" | grep -q "running"; then
            print_info "重启 $vm..."
            virsh reboot "$vm" # 发送重启信号
            if [ $? -eq 0 ]; then
                print_success "$vm 重启命令执行成功"
                restarted_any=true
            else
                print_error "$vm 重启失败"
            fi
        else
            print_warning "$vm 不在运行状态，尝试启动..."
            virsh start "$vm"
            restarted_any=true
        fi
    done
    
    if $restarted_any; then
        print_success "虚拟机重启命令已发送"
        echo ""
        
        # 等待一段时间后检查状态（重启需要更长时间）
        local restart_sleep_time=$((SLEEP_TIME + 10))
        wait_and_check_status $restart_sleep_time "重启"
        
        # 最终检查所有虚拟机是否都运行
        check_all_running
    else
        print_info "没有需要操作的虚拟机"
        show_status
    fi
}

# 显示用法
show_usage() {
    echo "用法: $0 {start|stop|restart|status|help}"
    echo ""
    echo "命令选项:"
    echo "  start   启动所有 Hadoop 虚拟机（等待 ${SLEEP_TIME} 秒后检查状态）"
    echo "  stop    关闭所有 Hadoop 虚拟机（等待 ${SLEEP_TIME} 秒后检查状态）"
    echo "  restart 重启所有 Hadoop 虚拟机（等待 $((SLEEP_TIME + 10)) 秒后检查状态）"
    echo "  status  查看所有 Hadoop 虚拟机状态"
    echo "  help    显示此帮助信息"
    echo ""
    echo "管理的虚拟机: ${VMs[*]}"
    echo "配置参数: 等待时间 = ${SLEEP_TIME} 秒"
}

# 主程序逻辑
case "$1" in
    start)
        start_vms
        ;;
    stop)
        shutdown_vms
        ;;
    restart)
        restart_vms
        ;;
    status)
        show_status
        ;;
    help|--help|-h)
        show_usage
        ;;
    *)
        print_error "错误: 未知命令 '$1'"
        echo ""
        show_usage
        exit 1
        ;;
esac

exit 0
